
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.security.SecureRandom;

import java.text.SimpleDateFormat;
import java.util.Date;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

public class AesCryption {
	public static boolean debug_mode = false;

	// 格式只到天，不然多个文件会有多个文件夹
	static SimpleDateFormat sdf = new SimpleDateFormat("yyyyMMdd");
	
	public static void main(String[] args) throws Exception {
		String input_file = "D:\\data\\title_text.txt";
		if (args.length > 0) {
			input_file = args.length > 0 ? args[0] : input_file;
		}

		// 共有2个参数 
		File in_file = new File(input_file);
		if (!in_file.exists()) {
			System.out.println("程序已退出，原因：The input file \"" + input_file + "\" does not exist.");
			return;
		}

		// 输出共包括3个文件：
		// 1. ./output/aes_20230416_162832/filename.ext.tmp
		// 2. ./output/aes_20230416_162832/filename.ext
		// 3. ./output/result.csv
		String out_dir = "./aes_" + sdf.format(new Date()) + "/";
		File folder = new File(out_dir);
		new File("./output/").mkdir();
        if (!folder.exists() && !folder.mkdir()) {
                System.out.println("程序已退出，原因：文件夹 " + folder + " 创建失败！");
				return;
		}

		String temp_file = out_dir + in_file.getName() + ".tmp";
		String output_file = out_dir + in_file.getName();

		if(debug_mode){
			System.out.print("------------------ Arguments Info -----------------\n");
			System.out.println("Input:    " + input_file);
			System.out.println("Tempfile: " + temp_file);
			System.out.println("Output:   " + output_file);
			System.out.print("---------------------------------------------------\n");
			return;
		}
			
		AesCryption ac = new AesCryption();
		for(int i = 0; i < 11; i ++){
			ac.process(input_file, temp_file, true);
		}
		Thread.sleep(100);
		for(int i = 0; i < 11; i ++){
			ac.process(temp_file, output_file, false);
		}
	}

	public static void fileTest() throws Exception {
		String input_file = "D:\\data\\title_text.txt";
		input_file = "D:\\data\\bts.bin";
		String encrypt_file = input_file + ".tmp";
		String output_file = input_file + input_file.substring(input_file.length() - 4);

		AesCryption ac = new AesCryption();
		ac.process(input_file, encrypt_file, true);
		ac.process(encrypt_file, output_file, false);

		System.out.println("Done.");
	}

	byte[] IV;
	SecretKey secKey;

	public AesCryption() {
		try {
			KeyGenerator keyGenerator = KeyGenerator.getInstance("AES");
			keyGenerator.init(128);
			secKey = keyGenerator.generateKey();
			IV = new byte[16];
			SecureRandom random = new SecureRandom();
			random.nextBytes(IV);
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public void process(String infile, String outfile, boolean encrypt) throws Exception {
		// 初始化所有相关变量
		FileInputStream fin = new FileInputStream(infile);
		FileOutputStream fout = new FileOutputStream(outfile);
		int filesize = fin.available();
		byte[] bts = new byte[encrypt ? 1024576 : 1024592];
		double t1 = 0, t2 = 0, t3 = 0;
		long timeTick = 0;
		int readCount = bts.length;
		Cipher cipher = getCipher(encrypt ? Cipher.ENCRYPT_MODE : Cipher.DECRYPT_MODE);

		// 加解密过程包括读取、加解密和写回文件。
		while (true) {
			// t1: 文件读取时间
			timeTick = System.nanoTime();
			readCount = fin.read(bts, 0, bts.length);
			if (readCount <= 0)
				break;
			t1 += System.nanoTime() - timeTick;

			// t2: 文件加密时间
			timeTick = System.nanoTime();
			if (readCount < bts.length) {
				byte[] bts1 = new byte[readCount];
				for (int i = 0; i < bts1.length; i++)
					bts1[i] = bts[i];
				bts = bts1;
			}
			byte[] result = cipher.doFinal(bts);
			t2 += System.nanoTime() - timeTick;

			// t3: 文件写入时间
			timeTick = System.nanoTime();
			fout.write(result, 0, result.length);
			t3 += System.nanoTime() - timeTick;
		}
		fin.close();
		fout.close();

		String datetime = new SimpleDateFormat("yyyy/MM/dd HH:mm:ss").format(new Date());
		int mode = encrypt ? 11 : 12;
		appendToFile(datetime, mode, "Read", t1, infile, filesize);
		appendToFile(datetime, mode, "Process", t2, infile, filesize);
		appendToFile(datetime, mode, "Write", t3, infile, filesize);
		appendToFile(datetime, mode, "Total", t1 + t2 + t3, infile, filesize);
	}

	// Cipher.ENCRYPT_MODE == 1
	// Cipher.DECRYPT_MODE == 2
	public Cipher getCipher(int type) throws Exception {
		Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
		SecretKeySpec keySpec = new SecretKeySpec(secKey.getEncoded(), "AES");
		IvParameterSpec ivSpec = new IvParameterSpec(IV);
		cipher.init(type, keySpec, ivSpec);
		return cipher;
	}


	public static boolean appendToFile(String datetime, int type, String item, double time, String infile, long size)
			throws IOException {
		String info = String.format("%s,%d,%s,%.6f,%.6f,%.2f,%s\n", // 格式化为统一的CSV字符串
				datetime, type, item, // 时间、类型（0=加密，1=解密）和项目名称
				time / 1000000000.0, // 输入时间单位为ns，转换为s
				size * 1000000000.0 / 1073741824.0 / time, size/1024576.0, infile);
		String resfile = "results_" + new SimpleDateFormat("yyyyMMdd").format(new Date()) + ".csv";
		if (!new File(resfile).exists())
			appendToFile(resfile, "DateTime,Type,Command,Time(s),Speed(GB/s),Flesize(B),Filepath\n");
		appendToFile(resfile, info);
		return true;
	}

	public static void appendToFile(String filepath, String content) throws IOException {
		System.out.print(content);
		FileWriter writer = new FileWriter(filepath, true);
		writer.write(content);
		writer.close();
	}
}
